import torch.nn.functional as F
from torch import nn

class LeNet5(nn.Module):
    def __init__(self):
        super(LeNet5, self).__init__()
        self.c1 = nn.Conv2d(1,6,5)  # input(1,32,32) output(6,28,28)
        self.c2 = nn.Conv2d(6,16,5) # input(6,14,14) output(16,10,10)
        self.pool = nn.MaxPool2d(2)
        self.fc1 = nn.Linear(16*5*5,120)
        self.fc2 = nn.Linear(120,84)
        self.fc3 = nn.Linear(84,10)

    def forward(self, x):
        x = F.relu(self.c1(x))   # (1,32,32)->(6,28,28)
        x = self.pool(x)            # (6,28,28)->(6,14,14)
        x = F.relu(self.c2(x))   # (6,14,14)->(16,10,10)
        x = self.pool(x)            # (16,10,10)->(16,5,5)
        x = x.view(-1,16*5*5)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x

